#pragma once
#ifndef CONV_H_
#define CONV_H_

#include "stream_tools.h"
#include "mbv2.h"

void MAC(ap_int<8> mul_a1, ap_int<8> mul_a2, ap_int<8> mul_b, ap_int<16>& rel1, ap_int<16>& rel2){
#pragma HLS INLINE
	ap_int<25> mul_a = ((ap_int<25>)mul_a1 << 16) + mul_a2;
	ap_int<32> result = mul_a * mul_b;
	rel1 = result(31,16) + result(15,15);
	rel2 = result(15,0);
}

template<
	unsigned N_IN,	// # of INput
	unsigned N_OUT,	// # of OUTput
	unsigned N_ICH,	// # of Input CHannel
	unsigned N_OCH,	// # of Output CHannel
	unsigned K,
	unsigned VEC_LEN
>
void conv(
	data_stream<N_IN * BIT_ACTV>& in, 
	data_stream<N_OUT * BIT_CONV>& out, 
	const ap_uint<N_OUT * N_IN * BIT_WGHT> weight[N_OCH / N_OUT][N_ICH / N_IN][K*K]
)
{
	static_assert(N_ICH >= N_IN,  "conv");
	static_assert(N_OCH >= N_OUT, "conv");
	static_assert(N_ICH % N_IN  == 0, "conv");
	static_assert(N_OCH % N_OUT == 0, "conv");

	constexpr unsigned FOLD_I = N_ICH / N_IN;
	constexpr unsigned FOLD_O = N_OCH / N_OUT;
	constexpr unsigned ITERS = VEC_LEN;

	assert(in.size() == VEC_LEN * FOLD_I * K * K);
	assert(out.empty());

	#pragma HLS bind_storage variable = weight type = rom_1p impl = lutram
	ap_uint<N_IN * BIT_ACTV> line[FOLD_I][K*K];
	ap_int<BIT_CONV> acc[N_OUT];
	#pragma HLS ARRAY_PARTITION variable=acc complete dim=1

	for (unsigned o = 0; o < N_OUT; ++o)
	{
		#pragma HLS UNROLL
		acc[o] = 0;
	}

	for (unsigned it = 0; it < ITERS; ++it)
	{
		for (unsigned fo = 0; fo < FOLD_O; ++fo)
		{
			for (unsigned fi = 0; fi < FOLD_I; ++fi)
			{
				for (unsigned k = 0; k < K*K; ++k)
				{
				#pragma HLS PIPELINE II=1
					// load
					ap_uint<N_IN * BIT_ACTV> in_buf;
					if (fo == 0)
					{
						in_buf = in.read();
						line[fi][k] = in_buf;
					}
					else
					{
						in_buf = line[fi][k];
					}
					ap_uint<N_OUT * N_IN * BIT_WGHT> wt_buf = weight[fo][fi][k];

					for (unsigned i = 0; i < N_IN; ++i)
					{
					#pragma HLS UNROLL
						ap_int<BIT_ACTV> x = in_buf(SLICE(BIT_ACTV, i));
						for (unsigned o = 0; o < N_OUT / 2; ++o)
						{
							ap_int<BIT_WGHT> w1 = wt_buf(SLICE(BIT_WGHT, N_IN * o * 2 + i));
							ap_int<BIT_WGHT> w2 = wt_buf(SLICE(BIT_WGHT, N_IN * o * 2 + N_IN + i));
							ap_int<16> tem1,tem2;
							MAC(w1, w2, x, tem1, tem2);
							acc[2 * o] += tem1;
							acc[2 * o + 1] += tem2;
						}
					}

					if(k == K*K -1 && fi == FOLD_I - 1)
					{
						ap_uint<N_OUT * BIT_CONV> out_buf;
						for (unsigned o = 0; o < N_OUT; ++o)
						{
						#pragma HLS UNROLL
							out_buf(SLICE(BIT_CONV, o)) = acc[o];
							acc[o] = 0;
						}
						out.write(out_buf);
					}
				}
			}

		}
	}

	assert(in.empty());
	assert(out.size() == VEC_LEN * FOLD_O);
	return;
};

template<
	unsigned N_IN,	// # of INput
	unsigned N_OUT,	// # of OUTput
	unsigned N_ICH,	// # of Input CHannel
	unsigned N_OCH,	// # of Output CHannel
	unsigned K,
	unsigned VEC_LEN
>
void conv_ek(
	data_stream<N_IN * K * BIT_ACTV>& in, 
	data_stream<N_OUT * BIT_CONV>& out,
	const ap_uint<N_OUT * N_IN * K * BIT_WGHT> weight[N_OCH / N_OUT][N_ICH / N_IN][K]
)
{
	static_assert(N_ICH >= N_IN,  "conv");
	static_assert(N_OCH >= N_OUT, "conv");
	static_assert(N_ICH % N_IN  == 0, "conv");
	static_assert(N_OCH % N_OUT == 0, "conv");

	constexpr unsigned FOLD_I = N_ICH / N_IN;
	constexpr unsigned FOLD_O = N_OCH / N_OUT;
	constexpr unsigned ITERS = VEC_LEN;

	assert(in.size() == VEC_LEN * FOLD_I * K);
	assert(out.empty());

	#pragma HLS bind_storage variable=weight type=rom_1p impl=lutram
	ap_uint<N_IN * K * BIT_ACTV> line[FOLD_I][K];
	#pragma HLS bind_storage variable=line type=ram_2p impl=lutram
	ap_int<BIT_CONV> acc[N_OUT];
	#pragma HLS array_partition variable=acc complete dim=1

	for (unsigned o = 0; o < N_OUT; ++o)
	{
		#pragma HLS UNROLL
		acc[o] = 0;
	}

	for (unsigned it = 0; it < ITERS; ++it)
	{
		for (unsigned fo = 0; fo < FOLD_O; ++fo)
		{
			for (unsigned fi = 0; fi < FOLD_I; ++fi)
			{
				for (unsigned k = 0; k < K; ++k)
				{
				#pragma HLS PIPELINE II=1
					// load
					ap_uint<N_IN * K * BIT_ACTV> in_buf;
					if (fo == 0)
					{
						in_buf = in.read();
						line[fi][k] = in_buf;
					}
					else
					{
						in_buf = line[fi][k];
					}
					ap_uint<N_OUT * N_IN * K * BIT_WGHT> wt_buf = weight[fo][fi][k];

					for (unsigned i = 0; i < N_IN * K; ++i)
					{
					#pragma HLS UNROLL
						ap_int<BIT_ACTV> x = in_buf(SLICE(BIT_ACTV, i));
						for (unsigned o = 0; o < N_OUT / 2; ++o)
						{
							ap_int<BIT_WGHT> w1 = wt_buf(SLICE(BIT_WGHT, N_IN * K * o * 2 + i));
							ap_int<BIT_WGHT> w2 = wt_buf(SLICE(BIT_WGHT, N_IN * K * o * 2 + N_IN * K + i));
							ap_int<16> tem1,tem2;
							MAC(w1, w2, x, tem1, tem2);
							acc[2 * o] += tem1;
							acc[2 * o + 1] += tem2;
						}
					}

					if(k == K - 1 && fi == FOLD_I - 1)
					{
						ap_uint<N_OUT * BIT_CONV> out_buf;
						for (unsigned o = 0; o < N_OUT; ++o)
						{
						#pragma HLS UNROLL
							out_buf(SLICE(BIT_CONV, o)) = acc[o];
							acc[o] = 0;
						}
						out.write(out_buf);
					}
				}
			}
		}
	}

	assert(in.empty());
	assert(out.size() == VEC_LEN * FOLD_O);
	return;
};

#endif
